package entity

import (
	"bufio"
	"encoding/json"
	"os"
)

type JsonlShardReader struct {
	file      *os.File
	scanner   *bufio.Scanner
	batchSize int
	first     bool
}

func NewJsonlShardReader(fileName string, batchSize int) (*JsonlShardReader, error) {
	file, err := os.Open(fileName)
	if err != nil {
		return nil, err
	}

	return &JsonlShardReader{
		file:      file,
		scanner:   bufio.NewScanner(file),
		batchSize: batchSize,
		first:     true,
	}, nil
}

func (r *JsonlShardReader) ReadBatch() ([]*Record, error) {
	count := 0
	var recordArray []*Record
	done := !r.scanner.Scan()
	for !done && count < r.batchSize {
		text := r.scanner.Text()
		var user User
		err := json.Unmarshal([]byte(text), &user)
		if err != nil {
			return nil, err
		}
		record := &Record{
			File:  r.file.Name(),
			User:  user,
			Start: false,
			Done:  false,
		}
		if r.first {
			record.Start = true
			r.first = false
		}
		recordArray = append(recordArray, record)
		count++
		done = !r.scanner.Scan()
	}
	// 记录文件完成的位置
	if done && count > 0 {
		recordArray[count-1].Done = true
	}
	if err := r.scanner.Err(); err != nil {
		return nil, err
	}
	return recordArray, nil
}

func (r *JsonlShardReader) Close() error {
	if r.file != nil {
		return r.file.Close()
	}
	return nil
}
